import sys
from pathlib import Path
PATH_ROOT = Path(__file__).parents[2]
sys.path.append(str(PATH_ROOT))

import cv2
import numpy as np
print("Loading YOLO modules (Loading Pytorch)...", flush=True, end='')
from ultralytics import YOLO
print("OK", flush=True)
import detection_gui.scripts.constants as const  # 仅保留必要的常量
import time

class YOLODetector:
    """YOLODetector 类用于加载指定的YOLO模型并进行螺丝检测任务"""
    def __init__(self, model_path: str = const.DEFECT_MODEL_PATH):  # 新模型路径
        assert Path(model_path).exists(), f"{model_path=} is not existed!"
        # 仅加载指定的检测模型
        self.model = YOLO(model_path, task='detect')
        self.conf_threshold = 0.7  # 固定置信度阈值为0.7
        self.avg_fps, self.count = 0.0, 0
    
    def __call__(self, img: np.ndarray, save: bool = False):
        start_time = time.time()
        assert img.ndim == 3, "Input image must be a 3D array (H, W, C)"
        if img.dtype == np.float32:
            img = (img * 255).astype(np.uint8)
        
        # 使用新模型进行检测，置信度固定为0.5
        results = self.model.predict(source=img, conf=self.conf_threshold, verbose=False)[0]

        detected_objects = []

        spot_num = 0
        wire_num = 0

        # 处理检测结果
        if results.boxes is not None:
            for i in range(len(results.boxes)):
                class_id = int(results.boxes[i].cls)
                class_name = self.model.names[class_id]
                bbox = results.boxes[i].xyxy[0].cpu().numpy().astype(int)
                conf = float(results.boxes[i].conf)
                if class_name == "spot":
                    spot_num += 1 
                elif class_name == "wire":
                    wire_num += 1
                detected_objects.append({
                    "name": class_name,
                    "conf": conf,
                    "bbox": bbox
                })

        # 计算符合置信度阈值（0.5）的物体数量
        object_count = len(detected_objects) - spot_num - wire_num

        # 可视化结果
        vis_img = self.visualize(img, detected_objects)
        if save:
            cv2.imwrite("prediction_result.jpg", vis_img)
            print("\n结果图片已保存为 prediction_result.jpg")
        
        # 计算FPS
        self.count += 1
        fps = 1 / (time.time() - start_time)
        self.avg_fps += (fps - self.avg_fps) / self.count

        # 返回结果包含物体数量
        return {
            "image": vis_img, 
            "detected_objects": detected_objects,
            "object_count": object_count,  # 新增：符合置信度的物体数量
            "spot_num": spot_num,
            "wire_num": wire_num,
            "avg_fps": self.avg_fps
        }

    def visualize(self, image: np.ndarray, objects: list):
        """在图上绘制检测结果"""
        vis_img = image.copy()
        
        # 绘制检测到的物体
        for obj in objects:
            bbox = obj['bbox']
            label = f"{obj['name']} ({obj['conf']:.2f})"
            color = const.DEFECT_COLORS[obj['name']]
            # 使用指定颜色绘制边界框
            cv2.rectangle(vis_img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
            cv2.putText(vis_img, label, (bbox[0], bbox[1] - 10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)

        return vis_img


if __name__ == '__main__':
    # --- 用户需要配置的参数 ---
    IMAGE_TO_PREDICT = '/root/Coding/ws_618/assets/gui.png'  # 请修改为你的图片路径
    # --- 执行预测 ---
    img = cv2.imread(IMAGE_TO_PREDICT)
    if img is None:
        print(f"无法读取图片: {IMAGE_TO_PREDICT}")
        sys.exit(1)
    
    # 初始化检测器
    yolo_detector = YOLODetector()
    result = yolo_detector(img, True)
    
    # 打印检测到的物体数量
    print(f"检测到置信度0.5以上的物体数量: {result['object_count']}")